# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/12 14:18 
@Author  : ZhangShenao 
@File    : 1.通过HuggingFaceHub调用开源模型.py
@Desc    : 通过HuggingFaceHub调用开源模型

"""
import dotenv
from langchain.chains.llm import LLMChain
from langchain_community.llms.huggingface_hub import HuggingFaceHub
from langchain_core.prompts import PromptTemplate

# 加载环境变量
dotenv.load_dotenv()

# 通过HuggingFaceHub,构建LLM
llm = HuggingFaceHub(
    repo_id='THUDM/chatglm3-6b',  # 指定模型仓库ID
)

# 构建Prompt
prompt = PromptTemplate.from_template('{query}')

# 构建LLM Chain
chain = LLMChain(llm=llm, prompt=prompt)

# 准备问题
query = 'Rose is which type of flower?'

# 调用Chain并返回结果
output = chain.run(query)

print(output)
